from trie_tree import TriedTree
import re
import random


def swap_word(new_words):
    random_idx_1 = random.randint(0, len(new_words) - 1)
    random_idx_2 = random_idx_1
    counter = 0
    while random_idx_2 == random_idx_1:
        random_idx_2 = random.randint(0, len(new_words) - 1)
        counter += 1
        if counter > 3:
            return new_words
    new_words[random_idx_1], new_words[random_idx_2] = " " + new_words[random_idx_2] + " ", " " + new_words[
        random_idx_1] + " "
    return new_words


class EntityHandler:
    def __init__(self, type_name):
        self.type_name = type_name
        self.segment = TriedTree()

        self.keyword_dic = {}
        self.re_dic = {}


class Entity:

    def __init__(self, content, e_type):
        self.content = content
        self.e_type = e_type

    def __str__(self):
        return "{}[{}]".format(self.content, self.e_type)

    def __eq__(self, other):
        if not isinstance(other, Entity):
            return False

        if self.content == other.content and self.e_type == other.e_type:
            return True

        return False


class Augment(object):
    GENERAL = 'O'

    def __init__(self):

        self.entity_handler_dic = {self.GENERAL: EntityHandler(self.GENERAL)}

    def add_keyword(self, keyword, text_list, entity_type=None):
        

        def do_add_keyword(type_name):
            self.entity_handler_dic.setdefault(
                type_name, EntityHandler(type_name))
            handler = self.entity_handler_dic.get(type_name)

            handler.keyword_dic.setdefault(keyword, [])

            keyword_texts = handler[keyword]

            handler.segment.insert(keyword, keyword_name)
            for text in text_list:

                for _ in re.finditer(keyword, text):

                    keyword_texts.append(keyword)


        if entity_type is None:
            do_add_entity(self.GENERAL)
            return
        assert isinstance(entity_type, list)
        for e_type in entity_type:
            do_add_entity(e_type)

    def add_regex_str(self, re_str, text_list, entity_type=None):
                
        def do_add_re_str(type_name):
            self.entity_handler_dic.setdefault(
                type_name, EntityHandler(type_name))
            handler = self.entity_handler_dic.get(type_name)

            source_texts = []
            source_texts.extend(text_list)

            handler.re_dic.setdefault(re_str, [])
            re_texts_exist = handler.re_dic.get(re_str)

            for text in source_texts:
                for item in re.finditer(re_str, text):
                    re_texts_exist.append(item.group())

        if entity_type is None:

            do_add_re_str(self.GENERAL)

            return

        for e_type in entity_type:

            do_add_re_str(e_type)

    def replace_text_by_regex(self, text, entity_type):
        handler = self.entity_handler_dic.get(entity_type)
        if handler is None:
            return text

        def get_regex_str_set():
            re_set = set()
            for re_key in handler.re_dic.keys():
                for _ in re.finditer(re_key, text):
                    re_set.add(re_key)
            return re_set

        text_list = list(text)
        result_text = text

        for re_str in get_regex_str_set():

            re_groups = list(re.finditer(re_str, result_text))

            re_groups.reverse()

            for item in re_groups:

                source_text = item.group()

                count = 0
                
                while True:

                    if count == 10:
                        break
                    aim_text = random.choice(handler.re_dic[re_str])

                    if aim_text == source_text:
                        count += 1
                        continue

                    text_list[item.start():item.end()] = aim_text
                    result_text = "".join(text_list)
                    count += 1
                    break

        return result_text

    def replace_text_by_entity(self, text, entity_type):

        handler = self.entity_handler_dic.get(entity_type)
        if handler is None:
            return text
        text_list = list(text)

        re_groups = list(handler.segment.segment_word(text))

        re_groups.reverse()

        result_text = text

        for item in re_groups:

            source_text = item["word"]
            keyword = item["tagger"]
            count = 0
            while True:
                if count == 3:
                    break
                aim_text = random.choice(handler.keyword_dic[keyword])
                if aim_text == source_text:
                    count += 1
                    continue

                text_list[item["begin"]:item["end"]] = aim_text
                result_text = "".join(text_list)

                break

        return result_text

    def random_swap(self, text, entity_type):
        words = self.seg_for_swap(text, entity_type)
        new_words = words.copy()
        new_words = swap_word(new_words)

        return "".join(new_words)

    def seg_for_swap(self, text, entity_type):
        handler = self.entity_handler_dic.get(entity_type)
        if handler is None:
            return [text]
        tokens = handler.segment.segment_word(text)

        if len(tokens) == 0:
            return [text]
        tokens.reverse()
        text_list = list(text)

        for tk in tokens:
            text_list[tk["begin"]:tk["end"]] = "##" + tk["word"] + "##"
        text = "".join(text_list)

        return text.split("##")

    def entity_augment(self, entity_list, n, rp=1, ep=1, sp=1):

        augment_list = []

        def is_same(new_entity_list):
            new_string = "".join([item.content for item in new_entity_list])
            string = "".join([item.content for item in entity_list])
            if new_string == string:
                return True

            return False

        for _ in range(n):
            temp_entity_list = []
            if random.random() < rp:

                for entity in entity_list:
                    aug_text = self.replace_text_by_regex(
                        entity.content, entity.e_type)
                    temp_entity_list.append(Entity(aug_text, entity.e_type))

                if not is_same(temp_entity_list):
                    augment_list.append(temp_entity_list)

            temp_entity_list = []
            if random.random() < ep:

                for entity in entity_list:
                    aug_text = self.replace_text_by_entity(
                        entity.content, entity.e_type)
                    temp_entity_list.append(Entity(aug_text, entity.e_type))
                if not is_same(temp_entity_list):
                    augment_list.append(temp_entity_list)

            temp_entity_list = []
            if random.random() < sp:

                for entity in entity_list:
                    aug_text = self.random_swap(entity.content, entity.e_type)
                    temp_entity_list.append(Entity(aug_text, entity.e_type))

                if not is_same(temp_entity_list):
                    augment_list.append(temp_entity_list)

        return augment_list

    def string_augment(self, text, n, rp=1, ep=1, sp=1):

        augment_list = []

        for _ in range(n):

            if random.random() < rp:
                temp_text = self.replace_text_by_regex(text, self.GENERAL)
                if temp_text != "":
                    augment_list.append(temp_text)
            if random.random() < ep:
                temp_text = self.replace_text_by_entity(text, self.GENERAL)
                if temp_text and temp_text != "":
                    augment_list.append(temp_text)
            if random.random() < sp:
                temp_text = self.random_swap(text, self.GENERAL)
                if temp_text and temp_text != "":
                    augment_list.append(temp_text)

        return augment_list

